import pandas as pd
import sys
from sklearn.model_selection import train_test_split  
from frame.training.ml_frame.month_rolling.model import *

class Ml_rolling_trin():
    def __init__(self, args) -> None:
        self.args = args

    def forard_process(self, X_train, X_eval, X_out, Y_train, Y_eval, Y_out, root, feature_list):
        x_train, _ = data_reshape(X_train)
        x_eval, _ = data_reshape(X_eval)
        x_out, _ = data_reshape(X_out)
        y_train, y_train_shape = data_reshape(Y_train)
        y_eval, y_eval_shape = data_reshape(Y_eval)
        y_out, y_out_shape = data_reshape(Y_out) 
        new_x_train, new_y_train = clean_data(x_train, y_train)
        new_x_eval, new_y_eval = clean_data(x_eval, y_eval)         
        x_out[np.isinf(x_out)] = 0
        x_train[np.isinf(x_train)] = 0
        x_eval[np.isinf(x_eval)] = 0
        x_out[np.isnan(x_out)] = 0
        x_train[np.isnan(x_train)] = 0
        x_eval[np.isnan(x_eval)] = 0        
        if self.args.model_name == 'lightgbm':
            ml_model = lgb_train( new_x_train, new_x_eval, new_y_train, new_y_eval, self.args, root)
        if self.args.model_name == 'catboost':
            ml_model = catboost_train( new_x_train, new_x_eval, new_y_train, new_y_eval, self.args)
        if self.args.model_name == 'liner': 
            ml_model = liner_train( new_x_train, new_x_eval, new_y_train, new_y_eval, self.args)
        if self.args.model_name == 'lasso':
            ml_model = lasso_train( new_x_train, new_x_eval, new_y_train, new_y_eval, self.args)  
        if self.args.model_name == 'random_forest':
            ml_model = random_forest_train( new_x_train, new_x_eval, new_y_train, new_y_eval, self.args)  
        out_pred = ml_model.predict(x_out)
        train_pred = ml_model.predict(x_train) 
        train_pred[np.isnan(y_train.squeeze(1))] = np.nan 
        eval_pred =  ml_model.predict(x_eval)
        eval_pred[np.isnan(y_eval.squeeze(1))] = np.nan 
        train_pred = data_reshape_inverse(train_pred, y_train_shape)
        eval_pred = data_reshape_inverse(eval_pred, y_eval_shape)
        out_pred = data_reshape_inverse(out_pred, y_out_shape)    
        return train_pred, Y_train, eval_pred, Y_eval, out_pred, Y_out


def clean_data(data_x, data_y):
    data_x[np.isinf(data_x)] = np.nan
    data_y[np.isinf(data_y)] = np.nan
    data = np.concatenate([data_x, data_y], axis=1)
    new_data = data[~np.isnan(data).any(axis=1), :]
    x = new_data[:,:-1]
    y = new_data[:,-1]
    return x, y
    
def data_reshape(array):
    shape = array.shape   
    re_array = array.reshape((-1, shape[-1]), order='F')
    return re_array, shape

def data_reshape_inverse(array, org_shape):
    re_array = array.reshape(org_shape, order='F')
    return re_array.squeeze(2)

def plot_lint_ml(array1, array2, name, root):
    plt.figure(figsize=(30, 6))
    plt.plot(array1,  alpha=0.5, label='pred')
    plt.plot(array2, alpha=0.5, label='true')
    plt.legend()
    save_root = f'{root}/{name}_result.png'
    plt.savefig(save_root)
    plt.close()
     
     
def train_eval_resplit(x_train, y_train, x_eval, y_eval):
    x = np.concatenate([x_train, x_eval], axis=0)
    y = np.concatenate([y_train, y_eval], axis=0)
    train_x, eval_x, train_y, eval_y = train_test_split( x, y, test_size=0.3, random_state=2022)
    return train_x, eval_x, train_y, eval_y